package cn.bugstack.xfg.dev.tech;

import cn.bugstack.xfg.dev.tech.config.OllamaConfig;
import com.alibaba.fastjson.JSON;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RunWith(SpringRunner.class)
@SpringBootTest
public class RAGTest {
    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private SimpleVectorStore simpleVectorStore;
    @Resource
    private PgVectorStore pgVectorStore;
    @Resource
    private OllamaChatClient chatClient;

    @Test
    public void upload() {

        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader("file.md");
        List<org.springframework.ai.document.Document> documents = tikaDocumentReader.get();
        List<Document> documentSplitterList = tokenTextSplitter.apply(documents);
        documents.forEach(document -> {
            log.info("---------------------documents:{}", document.getContent());
            document.getMetadata().put("knowledge", "知识库名称");
        });
        documentSplitterList.forEach(document -> {
            log.info("---------------------documentSplitterList:{}", document.getContent());
            document.getMetadata().put("knowledge", "知识库名称");
        });
        pgVectorStore.accept(documentSplitterList);
        log.info("上传完成");
    }

    @Test
    public void chat() {
        String message = "王大瓜，哪年出生";
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;
        // 指定文档搜索
        SearchRequest searchRequest = SearchRequest
                .query(message)
                .withTopK(5)
                .withFilterExpression("knowledge=='知识库名称'");
        //执行相似搜索
        List<Document> documents = pgVectorStore.similaritySearch(searchRequest);
        //合并搜索的结果
        String documentsText = documents.stream().map(Document::getContent).collect(Collectors.joining("\n"));
        //构建系统提示词

        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentsText));
        List<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message));
        messages.add(ragMessage);
        ChatResponse chatResponse = chatClient.call(
                new Prompt(messages, OllamaOptions.create().withModel("deepseek-r1:1.5b")

                ));
        log.info("chatResponse:{}", JSON.toJSONString(chatResponse));


    }
}
